#!/bin/bash
# ============================================================
# convert_mnist_to_trt.sh
# 一键将 MNIST ONNX 模型转换为 FP32 / FP16 的 TensorRT 引擎
# 打印当前环境 ONNX + TensorRT 版本信息
# ============================================================

set -e  # 遇到错误立即退出

# ===== 配置区域 =====
export model_name=mnist_model_epoch_01
export onnx_model=${model_name}.onnx
export timing_cache=${model_name}_timing_cache
export workspace_size=1024M  # 显存工作区大小，可根据 GPU 调整
# ====================

echo "============================================================"
echo "🧠 环境信息检查"
echo "------------------------------------------------------------"

# 打印 Python ONNX 版本
if command -v python3 &> /dev/null; then
    onnx_ver=$(python3 -c "import onnx; print(onnx.__version__)" 2>/dev/null || echo "未安装")
else
    onnx_ver="Python3 未安装"
fi
echo "🔹 ONNX version: $onnx_ver"

# 打印 TensorRT 版本
if command -v trtexec &> /dev/null; then
    trt_ver=$(trtexec --version | head -n 1)
else
    trt_ver="trtexec 未找到（TensorRT 可能未安装）"
fi
echo "🔹 TensorRT version: $trt_ver"

echo "============================================================"
echo "📁 模型信息"
echo "ONNX 模型路径: $onnx_model"
echo "============================================================"

# 检查模型是否存在
if [ ! -f "$onnx_model" ]; then
    echo "❌ 未找到 ONNX 模型文件: $onnx_model"
    exit 1
fi

# 检查 trtexec 是否可用
if ! command -v trtexec &> /dev/null; then
    echo "❌ 未找到 trtexec 工具，请确认 TensorRT 安装正确"
    exit 1
fi

# ===== FP32 转换 =====
echo "🚀 开始转换为 FP32 TensorRT 引擎..."
trtexec \
    --onnx=${onnx_model} \
    --saveEngine=${model_name}_fp32.trt \
    --timingCacheFile=${timing_cache} \
    --exportTimes=${model_name}_fp32_timing.json \
    --exportProfile=${model_name}_fp32_profile.json \
    --minShapes=input:1x1x28x28 \
    --optShapes=input:16x1x28x28 \
    --maxShapes=input:32x1x28x28 \
    --memPoolSize=workspace:${workspace_size} \
    --avgRuns=100 \
    --skipInference \
    2>&1 | tee ./build_${model_name}_fp32.log

echo "✅ FP32 TensorRT engine saved as: ${model_name}_fp32.trt"

# ===== FP16 转换 =====
echo "🚀 开始转换为 FP16 TensorRT 引擎..."
trtexec \
    --onnx=${onnx_model} \
    --saveEngine=${model_name}_fp16.trt \
    --timingCacheFile=${timing_cache} \
    --exportTimes=${model_name}_fp16_timing.json \
    --exportProfile=${model_name}_fp16_profile.json \
    --minShapes=input:1x1x28x28 \
    --optShapes=input:16x1x28x28 \
    --maxShapes=input:32x1x28x28 \
    --memPoolSize=workspace:${workspace_size} \
    --avgRuns=100 \
    --fp16 \
    --skipInference \
    2>&1 | tee ./build_${model_name}_fp16.log

echo "✅ FP16 TensorRT engine saved as: ${model_name}_fp16.trt"

# ===== 汇总结果 =====
echo "============================================================"
echo "🎉 所有转换完成！生成文件如下："
ls -lh ${model_name}_fp32.trt ${model_name}_fp16.trt
echo "============================================================"
